import asyncio
import concurrent.futures
import logging
import time

import aiohttp
from tqdm import tqdm

logging.basicConfig(level=logging.INFO)


# 异步发送请求的函数
async def send_request(session, url, method, headers, payload, progress_bar, is_async):
    try:
        # 发送POST或GET请求
        if method.upper() == 'POST':
            async with session.post(url, json=payload, headers=headers) as response:
                await handle_response(response, is_async, progress_bar)
        else:
            async with session.get(url, headers=headers) as response:
                await handle_response(response, is_async, progress_bar)


    except aiohttp.client_exceptions.ContentTypeError as e:
        logging.error(f"内容类型错误：{e}，处理 URL: {url}")
    except Exception as e:
        logging.error(f"发生了其他错误：{e}")
    finally:
        if not is_async:
            progress_bar.update(1)  # 更新进度条


# 处理响应的通用方法
async def handle_response(response, is_async, progress_bar):
    if is_async:
        progress_bar.update(1)  # 更新进度条
        return
    if response.status == 200:
        # 假设返回 JSON 数据
        if 'application/json' in response.headers.get('Content-Type', ''):
            await response.json()  # 解析 JSON 数据
        else:
            # 如果返回的是其他格式（例如 HTML），获取响应文本
            await response.text()
    else:
        res = await response.text()
        logging.error(f"错误：从 {response.url} 收到状态码 {response.status} 【{res}】")


# 线程工作函数
def thread_worker(url, method, headers, payload, num_requests, progress_bar, is_async):
    """异步"""
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)

    async def run_async_task():
        async with aiohttp.ClientSession() as session:
            tasks = [send_request(session, url, method, headers, payload, progress_bar, is_async) for _ in
                     range(num_requests)]
            await asyncio.gather(*tasks)

    loop.run_until_complete(run_async_task())
    loop.close()


# 进程工作函数
def process_worker(url, method, headers, payload, num_threads_per_process, num_requests_per_thread, is_async, config):
    """线程池"""
    total_requests = config['num_threads_per_process'] * config['num_requests_per_thread']
    progress_bar = tqdm(total=total_requests, desc="每个进程的线程 Requests", ncols=100, unit="次")

    with concurrent.futures.ThreadPoolExecutor(max_workers=num_threads_per_process) as executor:
        futures = [
            executor.submit(thread_worker, url, method, headers, payload, num_requests_per_thread, progress_bar,
                            is_async)
            for _ in range(num_threads_per_process)
        ]
        for future in futures:
            future.result()  # 等待线程任务完成


# 主函数
def main(config):
    """进程池"""
    start_time = time.time()
    with concurrent.futures.ProcessPoolExecutor(max_workers=config['num_processes']) as executor:
        futures = [
            executor.submit(process_worker, config['url'], config['method'], config['headers'], config['payload'],
                            config['num_threads_per_process'], config['num_requests_per_thread'], config['is_async'],
                            config)
            for _ in range(config['num_processes'])
        ]
        for future in futures:
            future.result()
    logging.info(f"················============ All requests completed in {time.time() - start_time:.2f} seconds.")


if __name__ == "__main__":
    config = {
        'url': 'https://baike.baidu.com/lemma/api/lemmapv?id=3b693b6bb391f939edde46b8',
        'headers': {
            'Content-Type': 'application/json',
            'User-Agent': 'Mozilla/5.0 (iPhone; CPU iPhone OS 16_6 like Mac OS X) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/16.6 Mobile/15E148 Safari/604.1 Edg/130.0.0.0'
        },
        'method': 'get',
        'payload': {},
        'num_processes': 1,  # 同时执行的进程数
        'num_threads_per_process': 1,  # 同时执行的线程数
        'num_requests_per_thread': 100000,  # 同时执行的异步任务
        'is_async': True
    }
    # 添加遍历次数添加任务
    main(config)
